// c-api-examples/speaker-identification-c-api.c
//
// Copyright (c)  2024  Xiaomi Corporation

// We assume you have pre-downloaded the speaker embedding extractor model
// from
// https://github.com/k2-fsa/sherpa-onnx/releases/tag/speaker-recongition-models
//
// An example command to download
// "3dspeaker_speech_campplus_sv_zh-cn_16k-common.onnx"
// is given below:
//
// clang-format off
//
// wget https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-recongition-models/3dspeaker_speech_campplus_sv_zh-cn_16k-common.onnx
//
// clang-format on
//
// Also, please download the test wave files from
//
// https://github.com/csukuangfj/sr-data

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "sherpa-onnx/c-api/c-api.h"

/**
 * @brief 计算输入音频文件的说话人嵌入向量（speaker embedding）
 * 
 * @param ex 说话人嵌入向量提取器的实例指针
 * @param wav_filename 输入的音频文件路径
 * @return const float* 返回说话人嵌入向量的指针（需由调用者负责释放内存）
 * 
 * @note 返回的嵌入向量指针v需要由调用者手动释放以避免内存泄漏
 */
static const float *ComputeEmbedding(
    const SherpaOnnxSpeakerEmbeddingExtractor *ex, const char *wav_filename) {
  // 1. 读取音频文件到内存
  const SherpaOnnxWave *wave = SherpaOnnxReadWave(wav_filename);
  if (wave == NULL) {
    // 读取失败处理
    fprintf(stderr, "Failed to read %s\n", wav_filename);
    exit(-1);
  }
 
  // 2. 创建在线流处理对象（用于实时处理音频流）
  const SherpaOnnxOnlineStream *stream =
      SherpaOnnxSpeakerEmbeddingExtractorCreateStream(ex);
 
  // 3. 将音频数据送入处理流
  //    参数：流对象、采样率、音频样本数据、样本数量
  SherpaOnnxOnlineStreamAcceptWaveform(stream, wave->sample_rate, wave->samples,
                                       wave->num_samples);
  
  // 4. 通知处理流音频输入已完成
  SherpaOnnxOnlineStreamInputFinished(stream);
 
  // 5. 检查音频是否足够长以生成有效的嵌入向量
  if (!SherpaOnnxSpeakerEmbeddingExtractorIsReady(ex, stream)) {
    fprintf(stderr, "The input wave file %s is too short!\n", wav_filename);
    exit(-1);
  }
 
  // 6. 计算说话人嵌入向量
  //    注意：返回的指针v需要由调用者负责释放
  const float *v =
      SherpaOnnxSpeakerEmbeddingExtractorComputeEmbedding(ex, stream);
 
  // 7. 释放临时资源
  SherpaOnnxDestroyOnlineStream(stream);  // 销毁流对象
  SherpaOnnxFreeWave(wave);               // 释放音频数据内存
 
  // 8. 返回嵌入向量（调用者需负责释放）
  return v;
}

int32_t main() {
  // 初始化说话人嵌入向量提取器的配置结构体
SherpaOnnxSpeakerEmbeddingExtractorConfig config;
 
// 将配置结构体内存清零（避免未初始化字段的随机值问题）
memset(&config, 0, sizeof(config));
 
// 设置模型路径（需提前下载模型文件）
// 模型下载地址：https://github.com/k2-fsa/sherpa-onnx/releases/tag/speaker-recognition-models
// 注意：这里使用的是中文普通话的3D Speaker模型（16kHz采样率）
config.model = "./3dspeaker_speech_campplus_sv_zh-cn_16k-common.onnx";
 
// 配置运行参数
config.num_threads = 1;    // 使用1个线程进行推理（可根据需要调整）
config.debug = 0;          // 关闭调试模式（0=关闭，1=开启）
config.provider = "cpu";   // 使用CPU进行推理（可选："cuda"使用GPU）
 
// 创建说话人嵌入向量提取器实例
const SherpaOnnxSpeakerEmbeddingExtractor *ex =
    SherpaOnnxCreateSpeakerEmbeddingExtractor(&config);
 
// 检查创建是否成功
if (!ex) {
    // 失败时打印错误信息并返回
    fprintf(stderr, "Failed to create speaker embedding extractor");
    return -1;
}

  // 获取说话人嵌入向量的维度（由模型决定）
int32_t dim = SherpaOnnxSpeakerEmbeddingExtractorDim(ex);
 
// 创建说话人嵌入向量管理器，用于管理多个说话人的注册和识别
// 参数dim指定嵌入向量的维度（由前面加载的模型决定）
const SherpaOnnxSpeakerEmbeddingManager *manager =
    SherpaOnnxCreateSpeakerEmbeddingManager(dim);
 
// 测试数据下载说明：
// 请从以下地址下载测试音频数据
// https://github.com/csukuangfj/sr-data
// 数据集包含多个说话人的注册/测试音频
 
// 定义说话人1的注册音频文件路径（3个不同音频用于注册）
const char *spk1_1 = "./sr-data/enroll/fangjun-sr-1.wav";
const char *spk1_2 = "./sr-data/enroll/fangjun-sr-2.wav";
const char *spk1_3 = "./sr-data/enroll/fangjun-sr-3.wav";
 
// 定义说话人2的注册音频文件路径（2个不同音频用于注册）
const char *spk2_1 = "./sr-data/enroll/leijun-sr-1.wav";
const char *spk2_2 = "./sr-data/enroll/leijun-sr-2.wav";
 
// 存储说话人1的嵌入向量数组（多注册音频）
// 注意：数组大小为4但只使用了3个元素（可能是为后续扩展预留）
const float *spk1_vec[4] = {NULL};  
// 计算每个注册音频的嵌入向量并保存
spk1_vec[0] = ComputeEmbedding(ex, spk1_1);
spk1_vec[1] = ComputeEmbedding(ex, spk1_2);
spk1_vec[2] = ComputeEmbedding(ex, spk1_3);
 
// 存储说话人2的嵌入向量数组（多注册音频）
const float *spk2_vec[3] = {NULL};
// 计算每个注册音频的嵌入向量并保存
spk2_vec[0] = ComputeEmbedding(ex, spk2_1);
spk2_vec[1] = ComputeEmbedding(ex, spk2_2);

  // 将说话人"fangjun"的多个嵌入向量注册到管理器中
// 参数：管理器实例、说话人名称、嵌入向量数组
// 注意：这里注册了之前计算的spk1_vec中的多个向量（多注册样本用于增强模型）
if (!SherpaOnnxSpeakerEmbeddingManagerAddList(manager, "fangjun", spk1_vec)) {
    fprintf(stderr, "Failed to register fangjun\n");  // 注册失败处理
    exit(-1);
}
 
// 验证说话人"fangjun"是否成功注册到系统中
if (!SherpaOnnxSpeakerEmbeddingManagerContains(manager, "fangjun")) {
    fprintf(stderr, "Failed to find fangjun\n");  // 未找到说话人处理
    exit(-1);
}
 
// 将说话人"leijun"的多个嵌入向量注册到管理器中
if (!SherpaOnnxSpeakerEmbeddingManagerAddList(manager, "leijun", spk2_vec)) {
    fprintf(stderr, "Failed to register leijun\n");  // 注册失败处理
    exit(-1);
}
 
// 验证说话人"leijun"是否成功注册到系统中
if (!SherpaOnnxSpeakerEmbeddingManagerContains(manager, "leijun")) {
    fprintf(stderr, "Failed to find leijun\n");  // 未找到说话人处理
    exit(-1);
}
 
// 验证系统中注册的说话人总数是否符合预期（应为2人）
if (SherpaOnnxSpeakerEmbeddingManagerNumSpeakers(manager) != 2) {
    fprintf(stderr, "There should be two speakers: fangjun and leijun\n");
    exit(-1);
}
 
// 获取所有已注册说话人的名称列表
const char *const *all_speakers = 
    SherpaOnnxSpeakerEmbeddingManagerGetAllSpeakers(manager);
 
// 打印所有已注册说话人列表
const char *const *p = all_speakers;
fprintf(stderr, "list of registered speakers\n-----\n");
while (p[0]) {  // 遍历以NULL结尾的字符串数组
    fprintf(stderr, "speaker: %s\n", p[0]);  // 打印每个说话人名称
    ++p;
}
fprintf(stderr, "----\n");
 
// 释放说话人列表内存（注意：不释放管理器本身）
SherpaOnnxSpeakerEmbeddingManagerFreeAllSpeakers(all_speakers);

  // 定义测试音频文件路径（包含已知和未知说话人的测试样本）
// test1: 已注册说话人fangjun的测试音频
// test2: 已注册说话人leijun的测试音频
// test3: 未注册说话人liudehua的测试音频（用于负样本测试）
const char *test1 = "./sr-data/test/fangjun-test-sr-1.wav";
const char *test2 = "./sr-data/test/leijun-test-sr-1.wav";
const char *test3 = "./sr-data/test/liudehua-test-sr-1.wav";
 
// 计算测试音频的嵌入向量（用于后续说话人识别）
const float *v1 = ComputeEmbedding(ex, test1);  // fangjun的测试样本
const float *v2 = ComputeEmbedding(ex, test2);  // leijun的测试样本
const float *v3 = ComputeEmbedding(ex, test3);  // liudehua的测试样本
 
// 设置识别阈值（相似度判断的门限值）
// 范围通常在0~1之间，值越大判断越严格
float threshold = 0.6;
 
// 测试1：识别fangjun的测试样本
// 在管理器中搜索与输入向量最匹配的说话人
// 参数：管理器实例、待识别向量、相似度阈值
const char *name1 = SherpaOnnxSpeakerEmbeddingManagerSearch(manager, v1, threshold);
if (name1) {
    // 找到匹配的说话人（相似度超过阈值）
    fprintf(stderr, "%s: Found %s\n", test1, name1);
    // 释放搜索结果字符串内存
    SherpaOnnxSpeakerEmbeddingManagerFreeSearch(name1);
} else {
    // 未找到匹配的说话人（所有相似度都低于阈值）
    fprintf(stderr, "%s: Not found\n", test1);
}
 
// 测试2：识别leijun的测试样本
const char *name2 = SherpaOnnxSpeakerEmbeddingManagerSearch(manager, v2, threshold);
if (name2) {
    fprintf(stderr, "%s: Found %s\n", test2, name2);
    SherpaOnnxSpeakerEmbeddingManagerFreeSearch(name2);
} else {
    fprintf(stderr, "%s: Not found\n", test2);
}
 
// 测试3：识别未注册说话人liudehua的测试样本（预期应不匹配）
const char *name3 = SherpaOnnxSpeakerEmbeddingManagerSearch(manager, v3, threshold);
if (name3) {
    // 理论上不应该进入这个分支（因为是未注册说话人）
    fprintf(stderr, "%s: Found %s\n", test3, name3);
    SherpaOnnxSpeakerEmbeddingManagerFreeSearch(name3);
} else {
    // 预期结果：未找到匹配的说话人
    fprintf(stderr, "%s: Not found\n", test3);
}

  // 验证测试：检查输入向量是否与指定说话人匹配
// 参数：管理器实例、目标说话人名称、待验证向量、相似度阈值
// 返回值：1表示匹配成功，0表示不匹配
int32_t ok = SherpaOnnxSpeakerEmbeddingManagerVerify(manager, "fangjun", v1, threshold);
if (ok) {
    // fangjun的测试样本应与fangjun注册信息匹配
    fprintf(stderr, "%s matches fangjun\n", test1);
} else {
    // 如果进入这里说明验证失败（可能是系统错误）
    fprintf(stderr, "%s does NOT match fangjun\n", test1);
}
 
// 验证leijun的测试样本是否与fangjun匹配（预期不匹配）
ok = SherpaOnnxSpeakerEmbeddingManagerVerify(manager, "fangjun", v2, threshold);
if (ok) {
    // 理论上不应该匹配（因为是不同说话人）
    fprintf(stderr, "%s matches fangjun\n", test2);
} else {
    // 预期结果：leijun的样本不应匹配fangjun
    fprintf(stderr, "%s does NOT match fangjun\n", test2);
}
 
// 测试说话人删除功能
fprintf(stderr, "Removing fangjun\n");
// 从管理器中删除说话人"fangjun"
if (!SherpaOnnxSpeakerEmbeddingManagerRemove(manager, "fangjun")) {
    // 删除失败处理
    fprintf(stderr, "Failed to remove fangjun\n");
    exit(-1);
}
 
// 验证删除后系统中剩余说话人数量
if (SherpaOnnxSpeakerEmbeddingManagerNumSpeakers(manager) != 1) {
    // 此时应只有leijun一个说话人（因为刚刚删除了fangjun）
    fprintf(stderr, "There should be only 1 speaker left\n");
    exit(-1);
}
 
// 尝试搜索已被删除的说话人fangjun的测试样本
name1 = SherpaOnnxSpeakerEmbeddingManagerSearch(manager, v1, threshold);
if (name1) {
    // 理论上不应该找到（因为对应说话人已被删除）
    fprintf(stderr, "%s: Found %s\n", test1, name1);
    SherpaOnnxSpeakerEmbeddingManagerFreeSearch(name1);
} else {
    // 预期结果：找不到匹配的说话人
    fprintf(stderr, "%s: Not found\n", test1);
}
 
// 删除最后一个说话人leijun
fprintf(stderr, "Removing leijun\n");
if (!SherpaOnnxSpeakerEmbeddingManagerRemove(manager, "leijun")) {
    // 删除失败处理
    fprintf(stderr, "Failed to remove leijun\n");
    exit(-1);
}
 
// 验证说话人数量是否为0（所有说话人都已删除）
if (SherpaOnnxSpeakerEmbeddingManagerNumSpeakers(manager) != 0) {
    // 错误处理：此时说话人数量应该为0
    fprintf(stderr, "There should be no speakers left\n");
    exit(-1);
}
 
// 尝试搜索leijun的测试样本（此时系统中已无任何说话人）
name2 = SherpaOnnxSpeakerEmbeddingManagerSearch(manager, v2, threshold);
if (name2) {
    // 理论上不应该找到（因为所有说话人都已被删除）
    fprintf(stderr, "%s: Found %s\n", test2, name2);
    SherpaOnnxSpeakerEmbeddingManagerFreeSearch(name2);
} else {
    // 预期结果：找不到任何匹配的说话人
    fprintf(stderr, "%s: Not found\n", test2);
}


  // 获取当前系统中所有已注册的说话人名称列表
// 返回一个以NULL结尾的字符串数组指针
char **all_speakers = SherpaOnnxSpeakerEmbeddingManagerGetAllSpeakers(manager);
 
// 遍历并打印所有已注册说话人
char **p = all_speakers;  // 临时指针用于遍历数组
fprintf(stderr, "list of registered speakers\n-----\n");
while (p[0]) {  // 当指针指向的字符串非空时继续循环
    fprintf(stderr, "speaker: %s\n", p[0]);  // 打印当前说话人名称
    ++p;  // 移动到下一个字符串指针
}
fprintf(stderr, "----\n");
 
// 释放说话人名称列表占用的内存
// 注意：需要使用配套的释放函数，不能直接使用free()
SherpaOnnxSpeakerEmbeddingManagerFreeAllSpeakers(all_speakers);
 
// 释放所有计算得到的嵌入向量内存
// 包括测试样本和注册样本的嵌入向量
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(v1);  // 释放test1的嵌入向量
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(v2);  // 释放test2的嵌入向量
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(v3);  // 释放test3的嵌入向量
 
// 释放说话人1的注册嵌入向量（共3个）
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(spk1_vec[0]);
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(spk1_vec[1]);
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(spk1_vec[2]);
 
// 释放说话人2的注册嵌入向量（共2个）
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(spk2_vec[0]);
SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(spk2_vec[1]);
 
// 销毁说话人嵌入管理器，释放所有相关资源
// 包括内部数据结构、注册的说话人模型等
SherpaOnnxDestroySpeakerEmbeddingManager(manager);
 
// 销毁说话人嵌入提取器，释放模型资源
SherpaOnnxDestroySpeakerEmbeddingExtractor(ex);
 
// 程序正常结束
return 0;
}
